/*
3D IoU Calculation and Rotated NMS(modified from 2D NMS written by others)
Written by Shaoshuai Shi
All Rights Reserved 2019-2020.
*/


#include <stdio.h>
#include <math.h>
#define THREADS_PER_BLOCK 16
#define CUDART_PI_F 3.141592654f
#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))

// #define DEBUG
const int THREADS_PER_BLOCK_NMS = sizeof(unsigned long long) * 8;
    
__device__ inline float grasp_distance(const float *grasp_a, const float *grasp_b, const float alpha, const float beta, const float theta){
    // params grasp_a: [x, y, z, av_x, av_y, av_z, cv_x, cv_y, cv_z, width]
    // alpha = 0.03, beta = 0.03, theta = 0.5
    float center_distance = sqrt((grasp_a[0]-grasp_b[0])*(grasp_a[0]-grasp_b[0]) + (grasp_a[1]-grasp_b[1])*(grasp_a[1]-grasp_b[1])+(grasp_a[2]-grasp_b[2])*(grasp_a[2]-grasp_b[2]));
    float approach_vector_dot = (grasp_a[3] * grasp_b[3] + grasp_a[4] * grasp_b[4] + grasp_a[5] * grasp_b[5]);
    if (approach_vector_dot > 1.0) {
        approach_vector_dot = 1.0;
    }
    else if (approach_vector_dot < -1.0) {
        approach_vector_dot = -1.0;
    }
    float approach_vector_distance = alpha * acos(approach_vector_dot)/CUDART_PI_F;
    float closing_vector_dot = fabs(grasp_a[6] * grasp_b[6] + grasp_a[7] * grasp_b[7] + grasp_a[8] * grasp_b[8]);
    if (closing_vector_dot > 1.0) {
        closing_vector_dot = 1.0;
    }
    float closing_vector_distance = beta * acos(closing_vector_dot)/CUDART_PI_F;
    float width_distance = theta * fabs(grasp_a[9] - grasp_b[9]);
    float distance = center_distance+approach_vector_distance+closing_vector_distance+width_distance;

    //#if __CUDA_ARCH__>=200
    //    printf("distance is : %f \n", distance);
    //    
    //#endif
    return distance;
}


__global__ void grasps_distance_kernel(const int num_a, const float *grasps_a, const int num_b, const float *grasps_b, float *ans_distance){
    // params grasps_a: (N, 7) [x, y, z, dx, dy, dz, heading]
    // params grasps_b: (M, 7) [x, y, z, dx, dy, dz, heading]
    const int a_idx = blockIdx.y * THREADS_PER_BLOCK + threadIdx.y;
    const int b_idx = blockIdx.x * THREADS_PER_BLOCK + threadIdx.x;

    if (a_idx >= num_a || b_idx >= num_b){
        return;
    }
    const float * cur_grasp_a = grasps_a + a_idx * 10;
    const float * cur_grasp_b = grasps_b + b_idx * 10;
    const float alpha = 0.03;
    const float beta  = 0.03;
    const float theta = 1.0;
    float s_distance = grasp_distance(cur_grasp_a, cur_grasp_b, alpha, beta, theta);
    ans_distance[a_idx * num_b + b_idx] = s_distance;
}


__global__ void grasp_nms_kernel(const int grasps_num, const float nms_metric_thresh,
                           const float *grasps, unsigned long long *mask){
    //params: grasps (N, 10) [x, y, z, dx, dy, dz, heading]
    //params: mask (N, N/THREADS_PER_BLOCK_NMS)

    //#if __CUDA_ARCH__>=200
    //    printf("thresh is : %f \n", nms_metric_thresh);
    //    
    //#endif

    const int row_start = blockIdx.y;
    const int col_start = blockIdx.x;

    // if (row_start > col_start) return;

    const int row_size = fminf(grasps_num - row_start * THREADS_PER_BLOCK_NMS, THREADS_PER_BLOCK_NMS);
    const int col_size = fminf(grasps_num - col_start * THREADS_PER_BLOCK_NMS, THREADS_PER_BLOCK_NMS);

    __shared__ float block_grasps[THREADS_PER_BLOCK_NMS * 10];

    if (threadIdx.x < col_size) {
        block_grasps[threadIdx.x * 10 + 0] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 0];
        block_grasps[threadIdx.x * 10 + 1] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 1];
        block_grasps[threadIdx.x * 10 + 2] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 2];
        block_grasps[threadIdx.x * 10 + 3] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 3];
        block_grasps[threadIdx.x * 10 + 4] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 4];
        block_grasps[threadIdx.x * 10 + 5] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 5];
        block_grasps[threadIdx.x * 10 + 6] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 6];
        block_grasps[threadIdx.x * 10 + 7] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 7];
        block_grasps[threadIdx.x * 10 + 8] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 8];
        block_grasps[threadIdx.x * 10 + 9] = grasps[(THREADS_PER_BLOCK_NMS * col_start + threadIdx.x) * 10 + 9];
    }
    __syncthreads();

    if (threadIdx.x < row_size) {
        const int cur_grasp_idx = THREADS_PER_BLOCK_NMS * row_start + threadIdx.x;
        const float *cur_grasp = grasps + cur_grasp_idx * 10;

        int i = 0;
        unsigned long long t = 0;
        const float alpha = 0.03;
        const float beta  = 0.03;
        const float theta = 1.0;
        int start = 0;
        if (row_start == col_start) {
          start = threadIdx.x + 1;
        }
        for (i = start; i < col_size; i++) {
            if (grasp_distance(cur_grasp, block_grasps + i * 10, alpha, beta, theta) < nms_metric_thresh){
                t |= 1ULL << i;
            }
        }
        const int col_blocks = DIVUP(grasps_num, THREADS_PER_BLOCK_NMS);
        mask[cur_grasp_idx * col_blocks + col_start] = t;
    }
}

void grasp_nmsLauncher(const float *grasps, unsigned long long * mask, int grasps_num, float nms_metric_thresh){
    dim3 blocks(DIVUP(grasps_num, THREADS_PER_BLOCK_NMS),
                DIVUP(grasps_num, THREADS_PER_BLOCK_NMS));
    dim3 threads(THREADS_PER_BLOCK_NMS);
    grasp_nms_kernel<<<blocks, threads>>>(grasps_num, nms_metric_thresh, grasps, mask);
}

void graspsdistanceLauncher(const int num_a, const float *grasps_a, const int num_b, const float *grasps_b, float *ans_distance){

    dim3 blocks(DIVUP(num_b, THREADS_PER_BLOCK), DIVUP(num_a, THREADS_PER_BLOCK));  // blockIdx.x(col), blockIdx.y(row)
    dim3 threads(THREADS_PER_BLOCK, THREADS_PER_BLOCK);

    grasps_distance_kernel<<<blocks, threads>>>(num_a, grasps_a, num_b, grasps_b, ans_distance);
#ifdef DEBUG
    cudaDeviceSynchronize();  // for using printf in kernel function
#endif
}
